﻿using System;
using System.Threading;
using System.Net;
using System.Net.Sockets;
using System.Threading.Tasks;

namespace AsynchronousMethods.ExtensibleNonblockingTPL
{
  public class Server
  {
    public void Serve (IPAddress address, int port)
    {
      ThreadPool.SetMinThreads (50, 50);
      TcpListener listener = new TcpListener (address, port);
      listener.Start ();
      while (true)
      {
        TcpClient c = listener.AcceptTcpClient ();
        ThreadPool.QueueUserWorkItem (DoReverseEcho, c);
      }
    }

    void DoReverseEcho (object client)
    {
      MessagingServices.BeginReverseEcho ((TcpClient)client, ReverseEchoCallback, client);
    }

    void ReverseEchoCallback (IAsyncResult r)
    {
      try
      {
        MessagingServices.EndReverseEcho (r);  // This will rethrow any exception
      }
      catch (Exception ex)
      {
        Console.WriteLine ("Error: " + ex.Message);
      }
      finally
      {
        ((TcpClient)r.AsyncState).Close ();
      }
    }
  }

  public class MessagingServices
  {
    public static IAsyncResult BeginReverseEcho (TcpClient client,
                                                 AsyncCallback callback,
                                                 object userState)
    {
      var re = new ReverseEcho ();
      re.Begin (client, callback, userState);
      return re;
    }

    public static byte [] EndReverseEcho (IAsyncResult r)
    {
      return ((ReverseEcho)r).End ();
    }
  }

  class ReverseEcho : IAsyncResult
  {
    TcpClient _client;
    NetworkStream _stream;
    object _userState;
    ManualResetEvent _waitHandle = new ManualResetEvent (false);
    int _bytesRead = 0;
    byte [] _data = new byte [5000];
    Exception _exception;

    internal ReverseEcho () { }

    // IAsyncResult members:

    public object AsyncState { get { return _userState; } }
    public WaitHandle AsyncWaitHandle { get { return _waitHandle; } }
    public bool CompletedSynchronously { get { return false; } }
    public bool IsCompleted
    {
      get { return _waitHandle.WaitOne (0, false); }
    }

    internal void Begin (TcpClient c, AsyncCallback callback, object state)
    {
      _client = c;
      _userState = state;
      _stream = _client.GetStream ();

      Task.Factory.StartNew (Read).ContinueWith (ant =>
      {
        _exception = ant.Exception;   // In case an exception occurred.

        if (_stream != null)
          try { _stream.Close (); }
          catch (Exception ex) { _exception = ex; };

        _waitHandle.Set ();

        if (callback != null) callback (this);
      });
    }

    internal byte [] End ()     // Wait for completion + rethrow any error.
    {
      AsyncWaitHandle.WaitOne ();
      if (_exception != null) throw _exception;
      return _data;
    }

    void Read ()
    {
      Task<int> readChunk = Task<int>.Factory.FromAsync (
        _stream.BeginRead, _stream.EndRead,
        _data, _bytesRead, _data.Length - _bytesRead, null);

      // NB: Note that we need TaskContinuationOptions.AttachedToParent to ensure that 
      // the parent task's cleanup continuations don't fire until we're finishing
      // writing the response. The first printing of the book is missing this!
      readChunk.ContinueWith (ContinueRead,
                              TaskContinuationOptions.NotOnFaulted
                            | TaskContinuationOptions.AttachedToParent);
    }

    void ContinueRead (Task<int> readChunk)
    {
      _bytesRead += readChunk.Result;
      if (readChunk.Result > 0 && _bytesRead < _data.Length)
      {
        Read ();       // More data to read!
        return;
      }
      Array.Reverse (_data);
      Task.Factory.FromAsync (_stream.BeginWrite, _stream.EndWrite,
                              _data, 0, _data.Length, null);
    }
  }
}
